#!/usr/bin/env python

import rospy
import actionlib
from geometry_msgs.msg import Twist, PoseStamped, Point
from move_base_msgs.msg import MoveBaseAction, MoveBaseGoal
from nav_msgs.msg import OccupancyGrid
from std_srvs.srv import Empty
from pynput import keyboard
import threading
import math
import numpy as np
import time
import tf

class TeleopWithExploration:
    def __init__(self):
        # 初始化ROS节点
        rospy.init_node('teleop_with_exploration', anonymous=True)
        
        # 创建Twist消息发布者
        self.cmd_vel_pub = rospy.Publisher('/cmd_vel', Twist, queue_size=10)
        
        # 导航客户端
        self.move_base_client = actionlib.SimpleActionClient('move_base', MoveBaseAction)
        
        # 地图订阅
        self.map_data = None
        self.map_resolution = 0.05  # 默认分辨率
        self.map_origin_x = 0
        self.map_origin_y = 0
        self.map_width = 0
        self.map_height = 0
        rospy.Subscriber('/map', OccupancyGrid, self.map_callback)
        
        # TF监听器
        self.tf_listener = tf.TransformListener()
        
        # 设置操作模式
        self.mode = "manual"  # "manual" 或 "auto"
        self.auto_nav_running = False
        self.nav_thread = None
        
        # 初始化速度值
        self.linear_speed = 0.2  # 线速度 (m/s)
        self.angular_speed = 1.0  # 角速度 (rad/s)
        self.twist = Twist()
        
        # 按键状态
        self.forward = False
        self.backward = False
        self.left = False
        self.right = False
        
        # 显示说明
        self.display_instructions()
        
        # 等待导航服务启动
        rospy.loginfo("等待move_base服务器...")
        try:
            self.move_base_client.wait_for_server(rospy.Duration(5.0))
            rospy.loginfo("move_base服务器已连接!")
        except:
            rospy.logwarn("无法连接到move_base服务器，自动导航功能可能不可用")

    def display_instructions(self):
        """显示操作说明"""
        rospy.loginfo("\n=== TurtleBot3 控制与SLAM探索系统 ===")
        rospy.loginfo("方向控制:")
        rospy.loginfo("  ↑ : 前进")
        rospy.loginfo("  ↓ : 后退")
        rospy.loginfo("  ← : 左转")
        rospy.loginfo("  → : 右转")
        rospy.loginfo("模式切换:")
        rospy.loginfo("  M : 手动模式")
        rospy.loginfo("  A : 自动探索模式")
        rospy.loginfo("  空格键 : 急停")
        rospy.loginfo("  ESC : 退出程序")
        rospy.loginfo("当前模式: 手动")
        rospy.loginfo("==========================\n")

    def map_callback(self, data):
        """接收地图数据"""
        self.map_data = data
        self.map_resolution = data.info.resolution
        self.map_origin_x = data.info.origin.position.x
        self.map_origin_y = data.info.origin.position.y
        self.map_width = data.info.width
        self.map_height = data.info.height
        
    def get_robot_position(self):
        """获取机器人当前位置"""
        try:
            (trans, rot) = self.tf_listener.lookupTransform('/map', '/base_footprint', rospy.Time(0))
            return trans[0], trans[1]
        except (tf.LookupException, tf.ConnectivityException, tf.ExtrapolationException):
            rospy.logwarn("无法获取机器人位置，使用(0,0)作为默认位置")
            return 0.0, 0.0

    def on_press(self, key):
        try:
            # 普通方向键处理
            if key == keyboard.Key.up:
                self.forward = True
            elif key == keyboard.Key.down:
                self.backward = True
            elif key == keyboard.Key.left:
                self.left = True
            elif key == keyboard.Key.right:
                self.right = True
            elif key == keyboard.Key.space:
                # 空格键 - 急停
                self.stop_robot()
                rospy.loginfo("急停!")
            elif key == keyboard.Key.esc:
                # ESC键 - 退出
                rospy.loginfo("程序退出!")
                return False
            elif hasattr(key, 'char'):
                if key.char in ['m', 'M']:
                    # 切换到手动模式
                    self.set_manual_mode()
                elif key.char in ['a', 'A']:
                    # 切换到自动导航模式
                    self.set_auto_mode()
        except AttributeError:
            pass
            
        if self.mode == "manual":
            self.update_speed()
        
        return True

    def on_release(self, key):
        try:
            # 释放键盘上的箭头键
            if key == keyboard.Key.up:
                self.forward = False
            elif key == keyboard.Key.down:
                self.backward = False
            elif key == keyboard.Key.left:
                self.left = False
            elif key == keyboard.Key.right:
                self.right = False
        except AttributeError:
            pass
            
        if self.mode == "manual":
            self.update_speed()

    def update_speed(self):
        """根据按键状态更新速度"""
        # 只在手动模式下更新速度
        if self.mode != "manual":
            return
            
        # 根据按键状态更新速度
        self.twist.linear.x = self.linear_speed if self.forward else (-self.linear_speed if self.backward else 0.0)
        self.twist.linear.y = 0.0
        self.twist.linear.z = 0.0
        self.twist.angular.x = 0.0
        self.twist.angular.y = 0.0
        self.twist.angular.z = self.angular_speed if self.left else (-self.angular_speed if self.right else 0.0)
        
        # 发布速度指令
        self.cmd_vel_pub.publish(self.twist)

    def set_manual_mode(self):
        """切换到手动模式"""
        if self.mode == "auto":
            # 如果当前是自动模式，则停止导航
            if self.auto_nav_running:
                self.stop_auto_navigation()
                
        self.mode = "manual"
        rospy.loginfo("已切换到手动模式")
        
        # 为安全起见，立即停止机器人
        self.stop_robot()

    def set_auto_mode(self):
        """切换到自动导航模式"""
        if self.mode == "manual":
            # 停止手动控制的任何动作
            self.stop_robot()
            
        self.mode = "auto"
        rospy.loginfo("已切换到自动探索模式")
        
        # 开始自动导航
        if not self.auto_nav_running:
            self.start_auto_navigation()

    def stop_robot(self):
        """发送停止命令"""
        twist = Twist()
        self.cmd_vel_pub.publish(twist)

    def start_auto_navigation(self):
        """启动自动导航线程"""
        if not self.auto_nav_running:
            self.auto_nav_running = True
            self.nav_thread = threading.Thread(target=self.auto_exploration_loop)
            self.nav_thread.daemon = True
            self.nav_thread.start()
            rospy.loginfo("开始自动探索")

    def stop_auto_navigation(self):
        """停止自动导航"""
        if self.auto_nav_running:
            self.auto_nav_running = False
            # 取消当前导航目标
            self.move_base_client.cancel_all_goals()
            rospy.loginfo("已停止自动导航")
            # 等待线程结束
            if self.nav_thread:
                self.nav_thread.join(timeout=1.0)

    def find_frontier_points(self):
        """寻找前沿点（未探索区域的边界）"""
        if self.map_data is None:
            rospy.logwarn("还没有接收到地图数据")
            return []
        
        # 获取地图数据
        width = self.map_width
        height = self.map_height
        data = np.array(self.map_data.data).reshape((height, width))
        
        # 定义未知区域、自由区域和障碍物
        unknown = -1
        free = 0
        occupied_threshold = 50
        
        # 找到自由区域与未知区域的边界
        frontiers = []
        for y in range(1, height-1):
            for x in range(1, width-1):
                # 检查当前点是否是自由区域
                if data[y, x] >= 0 and data[y, x] < occupied_threshold:
                    # 检查8个相邻点，看是否有未知区域
                    has_unknown_neighbor = False
                    for dy in [-1, 0, 1]:
                        for dx in [-1, 0, 1]:
                            if dx == 0 and dy == 0:
                                continue
                            if data[y+dy, x+dx] == unknown:
                                has_unknown_neighbor = True
                                break
                        if has_unknown_neighbor:
                            break
                    
                    # 如果是自由区域且至少有一个未知邻居，则是前沿点
                    if has_unknown_neighbor:
                        # 转换为世界坐标
                        world_x = self.map_origin_x + (x * self.map_resolution)
                        world_y = self.map_origin_y + (y * self.map_resolution)
                        frontiers.append((world_x, world_y))
        
        # 如果找到太多前沿点，进行下采样
        if len(frontiers) > 100:
            # 每隔n个点取一个
            n = len(frontiers) // 100
            frontiers = frontiers[::n]
            
        return frontiers

    def cluster_frontier_points(self, frontiers, cluster_tolerance=0.5):
        """将前沿点聚类，找出较大的未探索区域"""
        if not frontiers:
            return []
            
        # 简单聚类算法
        clusters = []
        for point in frontiers:
            # 检查该点是否属于现有的某个聚类
            assigned = False
            for i, cluster in enumerate(clusters):
                for cluster_point in cluster:
                    # 计算欧氏距离
                    dist = math.sqrt((point[0] - cluster_point[0])**2 + (point[1] - cluster_point[1])**2)
                    if dist < cluster_tolerance:
                        clusters[i].append(point)
                        assigned = True
                        break
                if assigned:
                    break
            
            # 如果不属于任何聚类，创建新的聚类
            if not assigned:
                clusters.append([point])
        
        # 过滤出较大的聚类（至少5个点）
        large_clusters = [c for c in clusters if len(c) >= 5]
        
        # 对每个聚类，计算中心点
        cluster_centers = []
        for cluster in large_clusters:
            x_sum = sum(p[0] for p in cluster)
            y_sum = sum(p[1] for p in cluster)
            center_x = x_sum / len(cluster)
            center_y = y_sum / len(cluster)
            cluster_centers.append((center_x, center_y))
            
        return cluster_centers

    def auto_exploration_loop(self):
        """自主探索的主循环"""
        try:
            # 等待接收到地图数据
            start_time = time.time()
            timeout = 10.0  # 10秒超时
            while self.map_data is None:
                if time.time() - start_time > timeout:
                    rospy.logerr("超时：未收到地图数据")
                    self.auto_nav_running = False
                    return
                rospy.sleep(0.5)
                
            rospy.loginfo("开始自动探索...")
            
            while self.auto_nav_running and not rospy.is_shutdown():
                # 寻找前沿点
                frontiers = self.find_frontier_points()
                
                if not frontiers:
                    rospy.loginfo("未找到探索前沿，尝试随机移动")
                    # 执行简单的旋转以查找未知区域
                    self.rotate_to_find_frontiers()
                    continue
                
                # 聚类前沿点并选择目标点
                cluster_centers = self.cluster_frontier_points(frontiers)
                
                if not cluster_centers:
                    rospy.loginfo("未找到有效的探索区域，尝试随机移动")
                    self.rotate_to_find_frontiers()
                    continue
                
                # 获取当前机器人位置
                robot_x, robot_y = self.get_robot_position()
                
                # 选择最近的聚类中心作为目标
                best_target = None
                min_distance = float('inf')
                for center in cluster_centers:
                    dist = math.sqrt((center[0] - robot_x)**2 + (center[1] - robot_y)**2)
                    if dist < min_distance:
                        min_distance = dist
                        best_target = center
                
                # 如果找到目标，设置导航目标
                if best_target:
                    x, y = best_target
                    rospy.loginfo(f"导航到探索前沿: x={x:.2f}, y={y:.2f}")
                    
                    # 设置导航目标
                    goal = MoveBaseGoal()
                    goal.target_pose.header.frame_id = "map"
                    goal.target_pose.header.stamp = rospy.Time.now()
                    goal.target_pose.pose.position.x = x
                    goal.target_pose.pose.position.y = y
                    goal.target_pose.pose.orientation.w = 1.0
                    
                    # 发送目标
                    self.move_base_client.send_goal(goal)
                    
                    # 等待直到达到目标、被取消或者超时
                    success = self.move_base_client.wait_for_result(rospy.Duration(30.0))
                    
                    if not self.auto_nav_running:
                        break
                        
                    state = self.move_base_client.get_state()
                    if state == actionlib.GoalStatus.SUCCEEDED:
                        rospy.loginfo("成功到达探索前沿")
                    else:
                        rospy.loginfo(f"无法到达探索前沿，状态: {state}")
                else:
                    rospy.loginfo("没有找到合适的探索目标")
                
                # 短暂暂停避免CPU占用过高
                rospy.sleep(2.0)
                
        except Exception as e:
            rospy.logerr(f"自动探索发生错误: {e}")
        finally:
            self.auto_nav_running = False

    def rotate_to_find_frontiers(self):
        """原地旋转以寻找新的前沿"""
        rospy.loginfo("执行原地旋转以寻找新区域")
        
        # 创建旋转命令
        twist = Twist()
        twist.angular.z = 0.5  # 缓慢旋转
        
        # 旋转大约10秒
        start_time = time.time()
        while time.time() - start_time < 10.0 and self.auto_nav_running:
            self.cmd_vel_pub.publish(twist)
            rospy.sleep(0.1)
            
        # 停止旋转
        twist.angular.z = 0.0
        self.cmd_vel_pub.publish(twist)

    def run(self):
        # 启动键盘监听
        with keyboard.Listener(on_press=self.on_press, on_release=self.on_release) as listener:
            rospy.loginfo("控制系统已启动，按ESC退出")
            listener.join()

if __name__ == '__main__':
    try:
        teleop = TeleopWithExploration()
        teleop.run()
    except rospy.ROSInterruptException:
        pass

